#include "include/main.h"
#include "include/process.h"
#include "include/rknn_api.h"
#include <cstdio>
#include <opencv2/core.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/opencv.hpp>
#include <set>
#include <string>

int main()
{
    std::string coco_class[80] = {"person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light",
           "fire hydrant","stop sign ","parking meter","bench","bird","cat","dog ","horse ","sheep","cow","elephant",
           "bear","zebra ","giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite",
           "baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle","wine glass","cup","fork","knife ",
           "spoon","bowl","banana","apple","sandwich","orange","broccoli","carrot","hot dog","pizza ","donut","cake","chair","sofa",
           "pottedplant","bed","diningtable","toilet ","tvmonitor","laptop","mouse","remote ","keyboard ","cell phone","microwave ",
           "oven ","toaster","sink","refrigerator ","book","clock","vase","scissors ","teddy bear ","hair drier", "toothbrush"};
    int model_len = 0;
    const int height = 1280;
    const int width = 1280;
    const int channel = 3;
    const char* model = "./model.rknn";
    const char* pic = "./bus.jpg";
    char* out_data = NULL;
    cv::Size size(1280, 1280);

    // 模型加载并返回模型二进制数据
    model_len = read_data_from_file(model, &out_data);
    if (out_data == NULL)
    {
        printf("load_model fail!\n");
        free(out_data);
        return -1;
    }

    // rknn 初始化
    rknn_context rknn_ctx = 0;
    int ret = rknn_init(&(rknn_ctx), out_data, model_len, 0);
    if (ret < 0)
    {
        printf("rknn_init fail! ret=%d\n", ret);
        free(out_data);
        return -1;
    }

    // 得到input output的num
    rknn_input_output_num io_num;
    ret = rknn_query(rknn_ctx, RKNN_QUERY_IN_OUT_NUM, &(io_num), sizeof(io_num));
    if (ret != RKNN_SUCC)
    {
        printf("rknn_query fail! ret=%d\n", ret);
        free(out_data);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    // Get Model Input Info
    rknn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (int i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        ret = rknn_query(rknn_ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC)
        {
            printf("rknn_query fail! ret=%d\n", ret);
            free(out_data);
            return -1;
        }
    }

    // Get Model Output Info
    rknn_tensor_attr output_attrs[io_num.n_output];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        ret = rknn_query(rknn_ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC)
        {
            printf("rknn_query fail! ret=%d\n", ret);
            free(out_data);
            return -1;
        }
    }

    // 读取jpg图像
    cv::Mat image = cv::imread(pic,cv::IMREAD_COLOR);
    cv::Mat img;
    cv::resize(image, img, size, cv::INTER_LINEAR);
    // set input
    rknn_input inputs[io_num.n_input];
    memset(inputs, 0, sizeof(inputs));
    for(int i = 0; i < io_num.n_input; i++)
    {
        inputs[i].index = i;
        inputs[i].type = RKNN_TENSOR_UINT8;
        inputs[i].fmt = RKNN_TENSOR_NHWC;
        inputs[i].size = channel*width*height;
        inputs[i].buf = img.data;
    }
    ret = rknn_inputs_set(rknn_ctx, io_num.n_input, inputs);
    if (ret < 0)
    {
        printf("rknn_input_set fail! ret=%d\n", ret);
        free(out_data);
        return -1;
    }

    // inference
    ret = rknn_run(rknn_ctx, nullptr);
    if (ret < 0)
    {
        printf("rknn_run fail! ret=%d\n", ret);
        free(out_data);
        return -1;
    }

    // get output
    rknn_output outputs[io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < io_num.n_output; i++)
    {
        outputs[i].index = i;
        outputs[i].want_float = false;
    }
    ret = rknn_outputs_get(rknn_ctx, io_num.n_output, outputs, NULL);
    if (ret < 0)
    {
        printf("rknn_outputs_get fail! ret=%d\n", ret);
        free(out_data);
        return -1;
    }

    // 数据后处理
    std::vector<float> filterBoxes;
    std::vector<float> objProbs;
    std::vector<int> classId;
    int validCount = 0;
    int stride = 0;
    int grid_h = 0;
    int grid_w = 0;
    int dfl_len = output_attrs[0].dims[2] / 4;
    int output_per_branch = io_num.n_output / 3;
    for (int i = 0; i < output_per_branch; i++)
    {
        void *score_sum = nullptr;
        int32_t score_sum_zp = 0;
        float score_sum_scale = 1.0;
        if (output_per_branch == 3){
            score_sum = outputs[i*output_per_branch + 2].buf;
            score_sum_zp = output_attrs[i*output_per_branch + 2].zp;
            score_sum_scale = output_attrs[i*output_per_branch + 2].scale;
        }
        int box_idx = i*output_per_branch;
        int score_idx = i*output_per_branch + 1;
        grid_h = output_attrs[box_idx].dims[1];
        grid_w = output_attrs[box_idx].dims[0];
        stride = height / grid_h;
        printf("tensordim:(%d,%d,%d,%d)\n",output_attrs[box_idx].dims[0],output_attrs[box_idx].dims[1],output_attrs[box_idx].dims[2],output_attrs[box_idx].dims[3]);
        validCount += process_u8((uint8_t *)outputs[box_idx].buf, output_attrs[box_idx].zp, output_attrs[box_idx].scale,
                                (uint8_t *)outputs[score_idx].buf, output_attrs[score_idx].zp, output_attrs[score_idx].scale,
                                (uint8_t *)score_sum, score_sum_zp, score_sum_scale,
                                grid_h, grid_w, stride, dfl_len,
                                filterBoxes, objProbs, classId, 0.20);
    }
    std::vector<int> indexArray;
    for (int i = 0; i < validCount; ++i)
    {
        indexArray.push_back(i);
    }
    std::set<int> class_set(std::begin(classId), std::end(classId));
    for (auto c : class_set)
    {
        nms(validCount, filterBoxes, classId, indexArray, c, 0.45);
    }

    // 画图
    for(int i = 0; i < validCount; i++)
    {
        if(indexArray[i] != -1)
        {
            cv::Rect rect(filterBoxes[i*4], filterBoxes[4*i+1]
                            ,filterBoxes[4*i+2], filterBoxes[4*i+3]);
            cv::rectangle(img, rect, cv::Scalar(255, 0, 0), 2);
            cv::putText(img, coco_class[classId[i]], cv::Point(filterBoxes[i*4], filterBoxes[4*i+1]),
                         cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255), 1);
            printf("classid:%d ",classId[i]);
            printf("%f ",objProbs[i]);
            printf("tensor:(%f,%f,%f,%f)\n",filterBoxes[4*i],filterBoxes[4*i+1],filterBoxes[4*i+2],filterBoxes[4*i+3]);
        }
    }

    cv::imwrite("./output.jpg", img);
    
    free(out_data);
}